{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "53224f5d-e050-4830-a95e-1ea0e50e1301",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "开始生成中国哲学八卦数据集...\n",
      "正在生成 乾卦 数据...\n",
      "乾卦 数据生成完成\n",
      "正在生成 坤卦 数据...\n",
      "坤卦 数据生成完成\n",
      "正在生成 震卦 数据...\n",
      "震卦 数据生成完成\n",
      "正在生成 巽卦 数据...\n",
      "巽卦 数据生成完成\n",
      "正在生成 坎卦 数据...\n",
      "坎卦 数据生成完成\n",
      "正在生成 离卦 数据...\n",
      "离卦 数据生成完成\n",
      "正在生成 艮卦 数据...\n",
      "艮卦 数据生成完成\n",
      "正在生成 兑卦 数据...\n",
      "兑卦 数据生成完成\n",
      "数据集已保存到 data/datasets/chinese_philosophy_trigrams_20250807.json\n",
      "数据集生成完成，保存位置: data/datasets/chinese_philosophy_trigrams_20250807.json\n"
     ]
    }
   ],
   "source": [
    "from zhipuai import ZhipuAI\n",
    "import json\n",
    "import os\n",
    "from datetime import datetime\n",
    "\n",
    "# 初始化智谱API客户端\n",
    "client = ZhipuAI(api_key=\"ef4c40ad6e90487c9cc38718c05e9480.UOR7nqzRAsBiBrPs\")  # 替换为您的API Key\n",
    "\n",
    "# 八卦基础信息\n",
    "eight_trigrams = [\n",
    "    {\"name\": \"乾\", \"symbol\": \"☰\", \"nature\": \"天\", \"position\": \"南\", \"family\": \"父\"},\n",
    "    {\"name\": \"坤\", \"symbol\": \"☷\", \"nature\": \"地\", \"position\": \"北\", \"family\": \"母\"},\n",
    "    {\"name\": \"震\", \"symbol\": \"☳\", \"nature\": \"雷\", \"position\": \"东北\", \"family\": \"长男\"},\n",
    "    {\"name\": \"巽\", \"symbol\": \"☴\", \"nature\": \"风\", \"position\": \"西南\", \"family\": \"长女\"},\n",
    "    {\"name\": \"坎\", \"symbol\": \"☵\", \"nature\": \"水\", \"position\": \"西\", \"family\": \"中男\"},\n",
    "    {\"name\": \"离\", \"symbol\": \"☲\", \"nature\": \"火\", \"position\": \"东\", \"family\": \"中女\"},\n",
    "    {\"name\": \"艮\", \"symbol\": \"☶\", \"nature\": \"山\", \"position\": \"西北\", \"family\": \"少男\"},\n",
    "    {\"name\": \"兑\", \"symbol\": \"☱\", \"nature\": \"泽\", \"position\": \"东南\", \"family\": \"少女\"}\n",
    "]\n",
    "\n",
    "def generate_trigram_philosophy(trigram):\n",
    "    \"\"\"使用GLM-4生成八卦的哲学解释\"\"\"\n",
    "    prompt = f\"\"\"\n",
    "    请从中国哲学角度详细解释八卦中的{trigram['name']}卦({trigram['symbol']})，包含以下内容：\n",
    "    1. 核心哲学思想（300字左右）\n",
    "    2. 道德寓意（200字左右）\n",
    "    3. 在人生决策中的应用（200字左右）\n",
    "    4. 与其他卦象的关系（100字左右）\n",
    "    使用中文回答，保持学术严谨性但避免过于晦涩。\n",
    "    \"\"\"\n",
    "\n",
    "    response = client.chat.completions.create(\n",
    "        model=\"glm-4\",\n",
    "        messages=[{\"role\": \"user\", \"content\": prompt}],\n",
    "        temperature=0.3,  # 降低随机性保证学术性\n",
    "        max_tokens=1500\n",
    "    )\n",
    "\n",
    "    return response.choices[0].message.content\n",
    "\n",
    "def parse_philosophy_text(text):\n",
    "    \"\"\"解析生成的哲学文本为结构化数据\"\"\"\n",
    "    sections = {\n",
    "        \"core_philosophy\": \"\",\n",
    "        \"moral_meaning\": \"\",\n",
    "        \"application\": \"\",\n",
    "        \"relationships\": \"\"\n",
    "    }\n",
    "\n",
    "    current_section = None\n",
    "    for line in text.split('\\n'):\n",
    "        if \"核心哲学思想\" in line:\n",
    "            current_section = \"core_philosophy\"\n",
    "        elif \"道德寓意\" in line:\n",
    "            current_section = \"moral_meaning\"\n",
    "        elif \"在人生决策中的应用\" in line:\n",
    "            current_section = \"application\"\n",
    "        elif \"与其他卦象的关系\" in line:\n",
    "            current_section = \"relationships\"\n",
    "        elif current_section and line.strip():\n",
    "            sections[current_section] += line.strip() + \"\\n\"\n",
    "\n",
    "    return sections\n",
    "\n",
    "def create_dataset():\n",
    "    \"\"\"创建完整数据集\"\"\"\n",
    "    dataset = {\n",
    "        \"metadata\": {\n",
    "            \"name\": \"中国哲学八卦数据集\",\n",
    "            \"version\": \"1.0\",\n",
    "            \"created_at\": datetime.now().isoformat(),\n",
    "            \"source\": \"智谱GLM-4生成\",\n",
    "            \"description\": \"包含易经八卦的哲学解释、象征意义和应用分析\"\n",
    "        },\n",
    "        \"trigrams\": []\n",
    "    }\n",
    "\n",
    "    for trigram in eight_trigrams:\n",
    "        print(f\"正在生成 {trigram['name']}卦 数据...\")\n",
    "        philosophy_text = generate_trigram_philosophy(trigram)\n",
    "        philosophy_data = parse_philosophy_text(philosophy_text)\n",
    "\n",
    "        trigram_data = {\n",
    "            \"basic_info\": trigram,\n",
    "            \"philosophy\": philosophy_data,\n",
    "            \"related_trigrams\": []\n",
    "        }\n",
    "\n",
    "        # 添加关联卦象关系\n",
    "        for other in eight_trigrams:\n",
    "            if other[\"name\"] != trigram[\"name\"]:\n",
    "                trigram_data[\"related_trigrams\"].append({\n",
    "                    \"name\": other[\"name\"],\n",
    "                    \"symbol\": other[\"symbol\"],\n",
    "                    \"relationship\": f\"与{trigram['name']}卦的关系待补充\"  # 实际应用中可进一步生成\n",
    "                })\n",
    "\n",
    "        dataset[\"trigrams\"].append(trigram_data)\n",
    "        print(f\"{trigram['name']}卦 数据生成完成\")\n",
    "\n",
    "    return dataset\n",
    "\n",
    "def save_dataset(dataset, format='json'):\n",
    "    \"\"\"保存数据集到文件\"\"\"\n",
    "    os.makedirs('data/datasets', exist_ok=True)\n",
    "    filename = f\"data/datasets/chinese_philosophy_trigrams_{datetime.now().strftime('%Y%m%d')}.{format}\"\n",
    "\n",
    "    if format == 'json':\n",
    "        with open(filename, 'w', encoding='utf-8') as f:\n",
    "            json.dump(dataset, f, ensure_ascii=False, indent=2)\n",
    "    # 可扩展其他格式如CSV\n",
    "\n",
    "    print(f\"数据集已保存到 {filename}\")\n",
    "    return filename\n",
    "\n",
    "# 主程序\n",
    "if __name__ == \"__main__\":\n",
    "    print(\"开始生成中国哲学八卦数据集...\")\n",
    "    dataset = create_dataset()\n",
    "    saved_file = save_dataset(dataset)\n",
    "    print(f\"数据集生成完成，保存位置: {saved_file}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "88c76ae8-1e36-497c-ad44-3da946bd0b9f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/miniconda3/envs/hlf_old_env/lib/python3.10/site-packages/huggingface_hub/file_download.py:945: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8a8ced257dbd44349f2968d036afd21c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/7 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "You are using an old version of the checkpointing format that is deprecated (We will also silently ignore `gradient_checkpointing_kwargs` in case you passed it).Please update to the new format on your modeling file. To use the new format, you need to completely remove the definition of the method `_set_gradient_checkpointing` in your model.\n",
      "You are using an old version of the checkpointing format that is deprecated (We will also silently ignore `gradient_checkpointing_kwargs` in case you passed it).Please update to the new format on your modeling file. To use the new format, you need to completely remove the definition of the method `_set_gradient_checkpointing` in your model.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "当前模型占用显存（单位为MiB）： 3739.6914672851562\n",
      "trainable params: 974,848 || all params: 6,244,558,848 || trainable%: 0.01561115883009451\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a418f164b51d412cb0c508c0acbebcc1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/7 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "比较微调前后的生成质量:\n",
      "\n",
      "问题 1: 什么是QLoRa技术？\n",
      "\n",
      "微调前模型生成:\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/miniconda3/envs/hlf_old_env/lib/python3.10/site-packages/transformers/generation/configuration_utils.py:392: UserWarning: `do_sample` is set to `False`. However, `temperature` is set to `0.7` -- this flag is only used in sample-based generation modes. You should set `do_sample=True` or unset `temperature`.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[gMASK]sop 问：什么是QLoRa技术？\n",
      "答：技术，如，如，如如等，如，如等，如，如，如，如，如，如，如，如，如，如，如，如，如，如，如，如，如，如，技术，如，技术，技术，数据，和，如，和，和，如，如，如，技术，如，如，技术，如，如，和，如，和，和，如，如，技术，如，如，如，技术，如，技术，数据，和，和，和，如，和，采集，如，数据，和，如，和，如，技术，如，功能，功能，功能，功能，功能，功能，功能，功能，功能，功能，功能，功能，功能，功能，功能，功能，功能，功能，功能，功能，功能，功能\n",
      "\n",
      "微调后模型生成:\n",
      "[gMASK]sop 问：什么是QLoRa技术？\n",
      "答：技术，如，如，如如等，如，如等，如，如，如，如，如，如，如，如，如，如，如，如，如，如，如，如，如，如，技术，如，技术，技术，数据，和，如，和，和，如，如，如，技术，如，如，技术，如，如，和，如，和，和，如，如，技术，如，如，如，技术，如，技术，数据，和，和，和，如，和，采集，如，数据，和，如，和，如，技术，如，功能，功能，功能，功能，功能，功能，功能，功能，功能，功能，功能，功能，功能，功能，功能，功能，功能，功能，功能，功能，功能，功能\n",
      "\n",
      "问题 2: 如何评估语言模型的性能？\n",
      "\n",
      "微调前模型生成:\n",
      "[gMASK]sop 问：如何评估语言模型的性能？\n",
      "答：性能评估的指标包括准确率，系统输出是否具有人类语言语言理解，语言生成，语言模型的语言理解，语言模型生成，语言模型的语言理解，语言生成，以及上下文生成等。为了获得更准确的评估，我们使用多个不同的评估指标进行综合评估。此外，为了评估语言生成模型的性能，我们通常会使用一些额外的指标，例如评估指标准确率，语言模型，准确率，如准确率，语言模型生成等。\n",
      "\n",
      "微调后模型生成:\n",
      "[gMASK]sop 问：如何评估语言模型的性能？\n",
      "答：性能评估的指标包括准确率，系统输出是否具有人类语言语言理解，语言生成，语言模型的语言理解，语言模型生成，语言模型的语言理解，语言生成，以及上下文生成等。为了获得更准确的评估，我们使用多个不同的评估指标进行综合评估。此外，为了评估语言生成模型的性能，我们通常会使用一些额外的指标，例如评估指标准确率，语言模型，准确率，如准确率，语言模型生成等。\n",
      "\n",
      "问题 3: Transformer架构的主要创新点是什么？\n",
      "\n",
      "微调前模型生成:\n",
      "[gMASK]sop 问：Transformer架构的主要创新点是什么？\n",
      "答：。等。等，等，等 Transformer架构，它采用自注意力机制，能够对整个上下文进行加权，从而实现对输入的语义表示。此外，该机制可以对不同位置的信息进行并行计算，从而提高计算效率。这种机制允许自适应地分配权重和计算，从而提高了整个系统的性能。\n",
      "\n",
      "微调后模型生成:\n",
      "[gMASK]sop 问：Transformer架构的主要创新点是什么？\n",
      "答：。等。等，等，等 Transformer架构，它采用自注意力机制，能够对整个上下文进行加权，从而实现对输入的语义表示。此外，该机制可以对不同位置的信息进行并行计算，从而提高计算效率。这种机制允许自适应地分配权重和计算，从而提高了整个系统的性能。\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "from langchain_community.chat_models import ChatZhipuAI\n",
    "from langchain.prompts import PromptTemplate\n",
    "from langchain.schema import HumanMessage\n",
    "from datasets import Dataset\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "import json\n",
    "import torch\n",
    "from transformers import AutoTokenizer, AutoModel, TrainingArguments, Trainer, BitsAndBytesConfig, AutoModelForCausalLM\n",
    "from peft import get_peft_model, LoraConfig, TaskType, prepare_model_for_kbit_training, PeftConfig, PeftModel\n",
    "from peft.utils import TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING\n",
    "\n",
    "target_modules = TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING['chatglm']\n",
    "lora_rank = 4                             # LoRA秩\n",
    "lora_alpha = 32                           # LoRA alpha值\n",
    "lora_dropout = 0.05                       # LoRA Dropout率\n",
    "\n",
    "# 修改数据生成函数\n",
    "def generate_synthetic_data(num_samples=1000):\n",
    "    # 首先检查文件是否存在\n",
    "    file_path = \"data/datasets/synthetic_data.json\"\n",
    "    if os.path.exists(file_path):\n",
    "        print(f\"检测到已有数据集文件 {file_path}，直接加载...\")\n",
    "        try:\n",
    "            with open(file_path, \"r\", encoding=\"utf-8\") as f:\n",
    "                existing_data = json.load(f)\n",
    "                print(f\"成功加载 {len(existing_data)} 条现有数据\")\n",
    "                return existing_data\n",
    "        except Exception as e:\n",
    "            print(f\"加载现有数据集失败，将重新生成。错误: {e}\")\n",
    "            \n",
    "    os.environ[\"ZHIPUAI_API_KEY\"] = \"ef4c40ad6e90487c9cc38718c05e9480.UOR7nqzRAsBiBrPs\"\n",
    "    \n",
    "    # 使用新的ChatZhipuAI类\n",
    "    # llm = ChatZhipuAI(model=\"chatglm3\", temperature=0.7)\n",
    "    llm = ChatZhipuAI(model=\"glm-4\", temperature=0.7)\n",
    "    # 提示词模板\n",
    "    prompt_template = PromptTemplate.from_template(\n",
    "        \"请生成一个关于{domain}的高质量问答对...\\n\"\n",
    "        \"输出格式为严格的JSON，只包含'question'和'answer'两个字段，\"\n",
    "        \"不要包含任何额外的文本、注释或代码标记。\"\n",
    "        \"确保answer字段中的内容不包含任何特殊控制字符。\"  # 明确要求纯净格式\n",
    "    )    \n",
    "    \n",
    "    domains = [\n",
    "        \"机器学习\",\n",
    "        \"深度学习\",\n",
    "        \"自然语言处理\",\n",
    "        \"计算机视觉\",\n",
    "        \"人工智能伦理\",\n",
    "        \"大数据\",\n",
    "        \"神经网络\",\n",
    "        \"强化学习\"\n",
    "    ]\n",
    "    \n",
    "    synthetic_data = []\n",
    "    \n",
    "    for _ in tqdm(range(num_samples), desc=\"生成合成数据\"):\n",
    "        domain = np.random.choice(domains)\n",
    "        prompt = prompt_template.format(domain=domain)\n",
    "        \n",
    "        try:\n",
    "            # 使用新的消息格式\n",
    "            response = llm([HumanMessage(content=prompt)])\n",
    "            response_content = response.content\n",
    "            # 处理响应内容（去除代码块标记（如果有））\n",
    "            if response_content.startswith(\"```json\"):\n",
    "                response_content = response_content[response_content.find(\"{\"):response_content.rfind(\"}\")+1]\n",
    "            # 过滤控制字符（ASCII码<32且不是换行符\\n）\n",
    "            response_content = ''.join(char for char in response_content if ord(char) >= 32 or ord(char) == 10)\n",
    "            data = json.loads(response_content)  # 尝试解释json\n",
    "            print(\"加载json：\", data)\n",
    "            if \"question\" in data and \"answer\" in data:  # 验证字段存在\n",
    "                synthetic_data.append(data)\n",
    "        except Exception as e:\n",
    "            print(f\"解析响应时出错: {e}\")\n",
    "            continue\n",
    "    \n",
    "    with open(\"data/datasets/synthetic_data.json\", \"w\", encoding=\"utf-8\") as f:\n",
    "        json.dump(synthetic_data, f, ensure_ascii=False, indent=2)\n",
    "    \n",
    "    return synthetic_data\n",
    "\n",
    "# 2. 准备数据集\n",
    "def prepare_dataset(data):\n",
    "    # 转换为HuggingFace数据集格式\n",
    "    dataset = Dataset.from_dict({\n",
    "        \"instruction\": [d[\"question\"] for d in data],\n",
    "        \"input\": [\"\"] * len(data),\n",
    "        \"output\": [d[\"answer\"] for d in data]\n",
    "    })\n",
    "    \n",
    "    # 分割训练集和验证集\n",
    "    dataset = dataset.train_test_split(test_size=0.1)\n",
    "    return dataset\n",
    "\n",
    "# 4. 数据预处理\n",
    "def preprocess_function(examples, tokenizer, max_length=512):\n",
    "    inputs = []\n",
    "    for inst, inp, outp in zip(examples[\"instruction\"], examples[\"input\"], examples[\"output\"]):\n",
    "        if inp:\n",
    "            text = f\"问：{inst}\\n输入：{inp}\\n答：\"\n",
    "        else:\n",
    "            text = f\"问：{inst}\\n答：\"\n",
    "        inputs.append(text)\n",
    "    \n",
    "    model_inputs = tokenizer(\n",
    "        inputs,\n",
    "        max_length=max_length,\n",
    "        truncation=True,\n",
    "        padding=\"max_length\"\n",
    "    )\n",
    "    \n",
    "    labels = tokenizer(\n",
    "        examples[\"output\"],\n",
    "        max_length=max_length,\n",
    "        truncation=True,\n",
    "        padding=\"max_length\"\n",
    "    )[\"input_ids\"]\n",
    "    \n",
    "    model_inputs[\"labels\"] = labels\n",
    "    return model_inputs\n",
    "\n",
    "def get_qlora_model(model_name_or_path):\n",
    "    _compute_dtype_map = {\n",
    "        'fp32': torch.float32,\n",
    "        'fp16': torch.float16,\n",
    "        'bf16': torch.bfloat16\n",
    "    }\n",
    "    \n",
    "    # QLoRA 量化配置\n",
    "    q_config = BitsAndBytesConfig(load_in_4bit=True,\n",
    "                                  bnb_4bit_quant_type='nf4',\n",
    "                                  bnb_4bit_use_double_quant=True,\n",
    "                                  bnb_4bit_compute_dtype=_compute_dtype_map['bf16'],\n",
    "                                  llm_int8_enable_fp32_cpu_offload=True  # 关键参数\n",
    "                                 )\n",
    "    # revision='b098244' 版本对应的 ChatGLM3-6B 设置 use_reentrant=False\n",
    "    # 最新版本 use_reentrant 被设置为 True，会增加不必要的显存开销\n",
    "    model = AutoModel.from_pretrained(model_name_or_path,\n",
    "                                      quantization_config=q_config,\n",
    "                                      device_map='auto',\n",
    "                                      use_cache=False,  # 训练时关闭（梯度检查点需要重新计算中间结果，而缓存会阻止这一过程）\n",
    "                                      trust_remote_code=True,\n",
    "                                      revision='b098244',\n",
    "                                      max_memory={0: \"7GiB\"}\n",
    "                                     )\n",
    "    \n",
    "    # 确保使用新版检查点\n",
    "    model.gradient_checkpointing_enable()\n",
    "\n",
    "    # 获取当前模型占用的 GPU显存（差值为预留给 PyTorch 的显存）\n",
    "    memory_footprint_bytes = model.get_memory_footprint()\n",
    "    memory_footprint_mib = memory_footprint_bytes / (1024 ** 2)  # 转换为 MiB\n",
    "    print(\"当前模型占用显存（单位为MiB）：\", memory_footprint_mib)\n",
    "    kbit_model = prepare_model_for_kbit_training(model)\n",
    "\n",
    "    # 创建一个LoraConfig对象，用于设置LoRA（Low-Rank Adaptation）的配置参数\n",
    "    lora_config = LoraConfig(\n",
    "        target_modules=target_modules,  # 指定将LoRA应用到的模型模块，通常是attention和全连接层的投影\n",
    "        r=lora_rank,                    # LoRA的秩，影响LoRA矩阵的大小\n",
    "        lora_alpha=lora_alpha,          # LoRA适应的比例因子\n",
    "        lora_dropout=lora_dropout,      # 在LoRA模块中使用的dropout率\n",
    "        bias='none',                    # 设置bias的使用方式，这里没有使用bias\n",
    "        inference_mode=False,           # 关闭推理模式\n",
    "        task_type=TaskType.CAUSAL_LM\n",
    "    )\n",
    "    # 使用get_peft_model函数和给定的配置来获取一个PEFT模型\n",
    "    qlora_model = get_peft_model(kbit_model, lora_config)  # 添加Lora\n",
    "    # 打印 QLoRA 微调训练的模型参数\n",
    "    qlora_model.print_trainable_parameters()\n",
    "    \n",
    "    return qlora_model\n",
    "\n",
    "\n",
    "# 6. 训练模型\n",
    "def train_model(qlora_model, tokenizer, dataset, model_name_or_path):\n",
    "    # 预处理数据集\n",
    "    column_names = dataset['train'].column_names\n",
    "    tokenized_dataset = dataset['train'].map(\n",
    "        lambda x: preprocess_function(x, tokenizer),\n",
    "        batched=True,\n",
    "        remove_columns=column_names\n",
    "    )\n",
    "   \n",
    "    # 定义训练参数（演示用，减少训练时间）\n",
    "    training_demo_args = TrainingArguments(\n",
    "        output_dir=f\"models/demo_0807/{model_name_or_path}\",  # 输出目录：保存模型和日志的路径\n",
    "        per_device_train_batch_size=8,                 # 每个GPU的训练batch大小：根据显存调整\n",
    "        gradient_accumulation_steps=4,                  # 梯度累积步数：模拟更大的batch size\n",
    "        learning_rate=1e-3,                             # 初始学习率：常用范围1e-5到1e-3\n",
    "        max_steps=100,                                  # 最大训练步数：控制演示时长\n",
    "        lr_scheduler_type=\"linear\",                     # 学习率调度策略：线性衰减\n",
    "        warmup_ratio=0.1,                               # 预热比例：前10%步数用于学习率预热\n",
    "        logging_steps=10,                               # 每10步记录一次日志\n",
    "        save_strategy=\"steps\",                          # 保存策略：按步数保存\n",
    "        save_steps=20,                                  # 每20步保存一次模型\n",
    "        optim=\"adamw_torch\",                            # 优化器：AdamW with torch实现\n",
    "        fp16=True,                                      # 启用混合精度训练：减少显存占用\n",
    "        report_to=[\"tensorboard\"],                      # 添加：记录指标到TensorBoard\n",
    "        gradient_checkpointing_kwargs={\"use_reentrant\": False},  # 显式设置\n",
    "    )\n",
    "    \n",
    "    # 初始化Trainer\n",
    "    trainer = Trainer(\n",
    "        model=qlora_model,                              # 要训练的模型（QLoRA适配后的模型）\n",
    "        args=training_demo_args,                        # 训练参数配置\n",
    "        train_dataset=tokenized_dataset,                # 训练数据集（已tokenize）\n",
    "        # data_collator=data_collator                    # 数据整理器：处理padding和batch\n",
    "        data_collator=lambda data: {\n",
    "            \"input_ids\": torch.stack([torch.tensor(d[\"input_ids\"]) for d in data]),\n",
    "            \"attention_mask\": torch.stack([torch.tensor(d[\"attention_mask\"]) for d in data]),\n",
    "            \"labels\": torch.stack([torch.tensor(d[\"labels\"]) for d in data])\n",
    "        }\n",
    "    )\n",
    "\n",
    "    # 开始训练并记录损失\n",
    "    print(\"开始训练...\")\n",
    "    train_result = trainer.train()                      # 返回的训练结果包含指标数据\n",
    "    print(\"train_result: \", train_result)\n",
    "\n",
    "    # 保存训练后的模型\n",
    "    trainer.model.save_pretrained(f\"models/demo0807/{model_name_or_path}\")\n",
    "    print(\"模型保存完成！\")\n",
    "    \n",
    "# 7. 比较微调前后的生成质量\n",
    "def compare_generation(model_before, model_after, tokenizer, test_questions):\n",
    "    print(\"\\n比较微调前后的生成质量:\")\n",
    "    \n",
    "    for i, question in enumerate(test_questions):\n",
    "        print(f\"\\n问题 {i+1}: {question}\")\n",
    "        \n",
    "        # 微调前生成\n",
    "        print(\"\\n微调前模型生成:\")\n",
    "        inputs = tokenizer(f\"问：{question}\\n答：\", return_tensors=\"pt\").to(model_before.device)\n",
    "        outputs = model_before.generate(**inputs, max_length=200, temperature=0.7)\n",
    "        print(tokenizer.decode(outputs[0], skip_special_tokens=True))\n",
    "        \n",
    "        # 微调后生成\n",
    "        print(\"\\n微调后模型生成:\")\n",
    "        inputs = tokenizer(f\"问：{question}\\n答：\", return_tensors=\"pt\").to(model_after.device)\n",
    "        outputs = model_after.generate(**inputs, max_length=200, temperature=0.7)\n",
    "        print(tokenizer.decode(outputs[0], skip_special_tokens=True))\n",
    "\n",
    "# 主流程\n",
    "def main():\n",
    "    # # 生成合成数据\n",
    "    # print(\"生成合成训练数据...\")\n",
    "    # synthetic_data = generate_synthetic_data(num_samples=500)  # 减少样本数以加快演示\n",
    "    \n",
    "    # # 准备数据集\n",
    "    # print(\"\\n准备数据集...\")\n",
    "    # dataset = prepare_dataset(synthetic_data)\n",
    "    # print(dataset)\n",
    "\n",
    "    # 获取组装了QloRa配置的模型\n",
    "    model_name_or_path = 'THUDM/chatglm3-6b'  # 模型ID或本地路径\n",
    "    qlora_model = get_qlora_model(model_name_or_path)\n",
    "    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path,\n",
    "                                          trust_remote_code=True,\n",
    "                                          use_reentrant=False,  # 推荐设置\n",
    "                                          revision='b098244')\n",
    "    \n",
    "    # print(\"\\n开始微调...\")\n",
    "    # model = train_model(qlora_model, tokenizer, dataset, model_name_or_path)\n",
    "    # print(\"\\n结束微调...\")\n",
    "\n",
    "    peft_model_path = 'models/demo0807/THUDM/chatglm3-6b'\n",
    "    compute_dtype = torch.bfloat16\n",
    "\n",
    "    # 2. 量化配置\n",
    "    q_config = BitsAndBytesConfig(\n",
    "        load_in_4bit=True,\n",
    "        bnb_4bit_quant_type=\"nf4\",\n",
    "        bnb_4bit_use_double_quant=True,\n",
    "        bnb_4bit_compute_dtype=compute_dtype,\n",
    "    )\n",
    "    \n",
    "    # 3. 加载基础模型\n",
    "    base_model = AutoModelForCausalLM.from_pretrained(\n",
    "        model_name_or_path,\n",
    "        quantization_config=q_config,\n",
    "        trust_remote_code=True,\n",
    "        device_map=\"auto\",\n",
    "        use_cache=True,\n",
    "        torch_dtype=compute_dtype,\n",
    "        revision='b098244'\n",
    "    )\n",
    "    base_model.requires_grad_(False)\n",
    "    base_model.eval()\n",
    "\n",
    "    # 5. 加载LoRA适配器\n",
    "    peft_config = PeftConfig.from_pretrained(peft_model_path)\n",
    "    peft_config.inference_mode = True\n",
    "    \n",
    "    model = PeftModel.from_pretrained(\n",
    "        base_model,\n",
    "        peft_model_path,\n",
    "        config=peft_config,\n",
    "        torch_dtype=compute_dtype\n",
    "    )\n",
    "    model.eval()\n",
    "  \n",
    "    # 比较生成质量\n",
    "    test_questions = [\n",
    "        \"什么是QLoRa技术？\",\n",
    "        \"如何评估语言模型的性能？\",\n",
    "        \"Transformer架构的主要创新点是什么？\"\n",
    "    ]\n",
    "    compare_generation(base_model, model, tokenizer, test_questions)\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    main()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e35842d2-a677-4216-b6d2-8d01287868d2",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "hlf_old_env",
   "language": "python",
   "name": "hlf_old_env"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
